#!/usr/bin/env python
# encoding: utf-8
"""
genDataGen.py

Created by Brant Faircloth on 2007-11-01.
Copyright (c) 2007 Michael J. Conroy and Brant C. Faircloth. All rights reserved.

Purpose
=======
Generate genetic data for a population over multiple generations.  Population allele frequencies
are initially generated from a Dirichlet given some alpha (1,...,1).  This is then fed into the 
multinomial for each "individual" we are simulating in the population.  2 draws are made from the
multinomial given the Dirichlet probabilities.  This becomes the genotype for an individual at a single
locus.  This is repeated over loci [1...n] given the number of alleles at each.  The initial Dirichlet
probs are retained, as are the allele frequencies resulting from the multinomial draw.  This becomes the
"starting" state for the population.  These genotypes are converted from quasi-ternary format to something 
more like a genotype.  The individuals are randomly assigned a sex maintaining a 50-50 sex ratio and 
(1 - survival) of them are "killed".  The remainder are allowed to breed, subject to the constraint that enough 
pairs bearing 5 offspring will replace the deaths in a population (yielding a lambda of 1).  These 
offspring are added back to the population, 70% of the pop. survives, &c., &c., &c.

Caveats
=======
No mutation, No linkage, Random Mating

Uses
=======
Testing software programs (relatedness, parentage, etc.).  Simulating data for a number of purposes.

"""

import numpy, pdb, random, sqlite3

popSize     = 100
loci        = [0,1,2,3,4,5]
sexRatio    = 0.5
alleles     = [2,5,10,7,12,6]
survival    = 0.50
generations = 2

def gen(alleles):
    '''generate genotype data given alleles and Dirichlet'''
    #pdb.set_trace()
    # set alpha for Dirichlet ([1 * alleles])
    alpha       = [1]*alleles
    # generate allele frequencies for *population* (not sample)
    prob        = numpy.ravel(numpy.random.dirichlet(alpha,1))
    for i in range(popSize):
        data    = numpy.ravel(numpy.random.multinomial(2,prob))
        try:
            genotype = numpy.vstack((genotype,data))
        except:
            genotype = data
    # keep track of counts and frequencies of alleles
    # at simulated locus
    count       = sum(genotype)
    frequency   = count/float(sum(count))
    return prob, genotype, count, frequency

def convert(genotypes):
    '''convert genotypes to readable format (allele number)'''
    #pdb.set_trace()
    for indiv in genotypes:
        # create an empty array where hstack will 'fail'
        # due to dimensions
        g = numpy.array([])
        for allele in range(len(indiv)):
            if indiv[allele] == 1:
                try:
                    g = numpy.hstack((g,allele+1))
                except:
                    g = numpy.array([allele+1])
            elif indiv[allele] == 2:
                g = numpy.array([allele+1,allele+1])
            else:
                pass
        try:
            pL = numpy.vstack((pL,g))
        except:
            pL = numpy.array([g])
        #print pL
    return pL

def rsample(n, r):
    "Generate r randomly chosen, sorted integers from [0,n] w/o replacement"
    r   = int(r)
    s   = random.sample(xrange(n),r)
    return s

def sex(g):
    '''determine sex of simulated genotypes'''
    # determine numbers of each sex
    n = int(len(g) * sexRatio)
    # randomly sample who will be male/female
    s = rsample(len(g),n)
    #s = list(rsample(len(g),n))
    # 1 = male
    # 0 = female
    sList = []
    # apply sexes to the indices
    for i in range(len(g)):
        if i in s:
            sList.append([1])
        else:
            sList.append([0])
    sex = numpy.array(sList)
    # stack these sexes onto the genotype array
    gSex = numpy.column_stack((sex,g))
    #pdb.set_trace()
    # return genetic data with sex appended
    return gSex

def iden(g,idS):
    '''assign id numbers to individuals'''
    # pick some sequential index numbers given start position
    id = numpy.arange(idS,idS+len(g))
    # stack them onto the genotypes
    g = numpy.column_stack((id,g))
    # add the end of the index to the start value
    # so that ids do not overlap (i,e. they are unique)
    idS = id[-1]+1
    #pdb.set_trace()
    # return genetic data with identification appended as 1st column
    return g, idS

def surviv(g):
    '''remove some subset of individuals from the population'''
    # determine popSize to keep
    keep = int(len(g) * survival)
    # sample individuals to keep (or randomly select
    # removals...)
    s = rsample(len(g),keep)
    # sort sample to keep so that ids are in proper order
    s.sort()
    # reindex the genotypes minus the removals
    s = numpy.array(s)
    surv = g[s]
    # return an array of individuals that survived
    return surv

def breed(n,g,survival,idStart,offspring=5):
    '''allow remainder to randomly breed such that lambda = 1 (e.g. birth = deaths)'''
    # set number of breeders * # offspring = deaths (lambda = 1)
    breeders = ((1-survival) * n)/offspring
    m   = g[:,1]==1
    f   = g[:,1]==0
    # create arrays of males and females from above
    males   = g[m]
    females = g[f]
    # randomly select number of breeders of each sex
    mI  = rsample(len(males),breeders)
    fI  = rsample(len(females),breeders)
    # since list is sorted, randomly permute it
    # and return randomly sorted array rather
    # than adjust the rsample function
    mI  = numpy.random.permutation(mI)
    fI  = numpy.random.permutation(fI)
    # get the data for the randomly selected individuals
    mIG,fIG = males[mI],females[fI]
    # create list for parents
    parents = []
    # loop through the individuals
    for i in range(len(mIG)):
        # loop through the alleles
        genotypeLength = range(0,len(mIG[i][2:]),2)
        for k in range(offspring):
            # input parental ids
            parents.append([mIG[i][0],mIG[i][1],fIG[i][0],fIG[i][1]])
            for j in genotypeLength:
                if j < genotypeLength[-1]:
                    mLoc = mIG[i][2:][j:j+2]
                    fLoc = fIG[i][2:][j:j+2]
                else:
                    mLoc = mIG[i][2:][j:]
                    fLoc = fIG[i][2:][j:]
                # run binomial trial w/ p = 0.5 for mendelian
                # inheritance at each locus
                mBinomial,fBinomial = numpy.random.binomial(1,0.5),numpy.random.binomial(1,0.5)
                jLoc = [mLoc[mBinomial],fLoc[fBinomial]]
                jLoc.sort()
                try:
                    jIG = numpy.hstack((jIG,jLoc))
                except:
                    jIG = numpy.array(jLoc)
            try:
                jG = numpy.vstack((jG,jIG))
            except:
                jG = jIG
            jIG = numpy.array([])
    # give offspring sex
    jG = sex(jG)
    # give offspring ids
    jG,idStart = iden(jG,idStart)
    # make parents an array
    parents = numpy.array(parents)
    # concat parent ids with offspring ids
    parentage = numpy.column_stack((parents,jG[:,0:2]))
    # append juveniles to the bottom of the genotype array
    g = numpy.vstack((g,jG))
    #pdb.set_trace()
    return g,parentage,idStart

def alleleFreqs(g):
    freqs = {}
    for l in range(len(loci)):
        freqs[l]= numpy.zeros(alleles[l])
    # for all individuals in the genotype array    
    #pdb.set_trace()
    for indiv in g:
        genotypeLength = range(2,len(indiv),2)
        for genotype in genotypeLength:
            if genotype < genotypeLength[-1]:
                locus = indiv[genotype:genotype+2]
            else:
                locus = indiv[genotype:]
            position = genotypeLength.index(genotype)
            for allele in locus:
                freqs[position][int(allele-1)]+=1
    #pdb.set_trace()
    # get population allele frequencies
    keys = freqs.keys()
    keys.sort()
    for i in keys:
        freqs[i]=(freqs[i]/sum(freqs[i]))
    return freqs

def database(c, number):
    connect.execute('''create table paren_?''', number)
    connect.execute('''create table geno_?''', number)
    connect.execute('''create table freq_?''', number)
    cconnect.commit()
    
def main():
    #connect = sqlite3.connect('parentageSimulation')
    # staring identification number
    idStart = 1
    prob,freqs,counts = {},{},{}
    pG = numpy.array([])
    for l in range(len(loci)):
        p,g,c,f     = gen(alleles[l])
        prob[l]     = numpy.round(p,3)
        freqs[l]    = f
        counts[l]   = c
        # prettify the genotypes
        pL      = convert(g)
        try:
            pG  = numpy.column_stack((pG,pL))
        except:
            pG  = pL
    #pdb.set_trace()
    pG          = sex(pG)
    pG,idStart  = iden(pG, idStart)
    for i in range(generations):
        pG                      = surviv(pG)
        pG,parentage,idStart    = breed(popSize,pG,survival,idStart)
        gProb                   = alleleFreqs(pG)
        #pdb.set_trace()
        print pG
        print parentage
        if (generations - i) < 3:
            #database(connect,i)
            genotypes_file = ('genotypes_gen_%s') % (i)
            #pdb.set_trace()
            file = open(genotypes_file,'w')
            file.write('numpy.array([')
            for genotype in pG:
                file.write('[')
                for element in genotype:
                    file.write(('%s, ') % (int(element)))
                file.write('],\n')
            file.close()
            parentage_file = ('parentage_gen_%s') % (i)
            file = open(parentage_file,'w')
            file.write(('%s\n') % (gProb))
            for parents in parentage:
                file.write(('%s\n') % (parents))
            file.close()
    pdb.set_trace()

if __name__ == '__main__':
    main()